#不要动这行注释一个字！ 文件名： app_ui.py 每次修改后版本号+0.1
#当前版本号：3.3.1
import tkinter as tk
from tkinter import filedialog, messagebox
from PIL import Image, ImageTk
import threading
import os
import traceback
import time
import torch
import numpy as np
from PIL import Image
from tqdm.auto import tqdm
from skimage.color import rgb2lab, lab2rgb
import cv2

try:
    import slic_torch
    algo_loaded = True
except Exception as e:
    algo_loaded = False
    algo_load_error = e
    print(f"算法模块加载失败: {e}")

class PixelArtApp:
    _instance = None

    def __new__(cls, *args, **kwargs):
        if cls._instance is not None:
            try:
                cls._instance.root.destroy()
            except Exception:
                pass
        cls._instance = super().__new__(cls)
        return cls._instance

    def __init__(self, root):
        if hasattr(self, 'initialized') and self.initialized:
            return
        self.initialized = True

        self.root = root
        self.root.title("PyTorch SLIC 像素画生成器")
        self.root.geometry("900x900")

        self.orig_w = None
        self.orig_h = None
        self.updating_fields = False
        self.update_job = None

        # 参数绑定
        self.target_width_str = tk.StringVar(value="64")
        self.target_height_str = tk.StringVar(value="")
        self.region_size_str = tk.StringVar(value="16")

        self.color_weight = tk.DoubleVar(value=10.0)
        self.iter_num = tk.IntVar(value=10)
        self.max_color_number = tk.IntVar(value=16)

        self.enable_pixelization = tk.IntVar(value=0)
        self.enable_slic = tk.IntVar(value=1)

        self.compress_mode = tk.StringVar(value="pixel")  # none/pixel/interpolate

        self.noise_filter_strength = tk.IntVar(value=0)

        self.image_path = None
        self.folder_path = None

        # 文件夹实时预览
        self.processed_images = []
        self.current_index = 0

        self.target_width_str.trace_add("write", self.schedule_auto_update)
        self.target_height_str.trace_add("write", self.schedule_auto_update)
        self.region_size_str.trace_add("write", self.schedule_auto_update)

        self.create_ui()

        # 绑定方向键切换
        self.root.bind("<Left>", self.show_prev_image)
        self.root.bind("<Right>", self.show_next_image)

    def create_ui(self):
        frame = tk.Frame(self.root)
        frame.pack(side="left", fill="y", padx=10, pady=10)

        tk.Label(frame, text="目标宽度 target_width:").pack(anchor="w")
        tk.Entry(frame, textvariable=self.target_width_str).pack(fill="x", pady=3)

        tk.Label(frame, text="目标高度 target_height (留空自动计算):").pack(anchor="w")
        tk.Entry(frame, textvariable=self.target_height_str).pack(fill="x", pady=3)

        tk.Label(frame, text="超像素大小 region_size:").pack(anchor="w")
        tk.Entry(frame, textvariable=self.region_size_str).pack(fill="x", pady=3)

        tk.Label(frame, text="颜色权重 color_weight:").pack(anchor="w")
        tk.Entry(frame, textvariable=self.color_weight).pack(fill="x", pady=3)

        tk.Label(frame, text="迭代次数 iter_num:").pack(anchor="w")
        tk.Entry(frame, textvariable=self.iter_num).pack(fill="x", pady=3)

        tk.Label(frame, text="压缩/像素化模式:").pack(anchor="w")
        tk.OptionMenu(frame, self.compress_mode, "none", "pixel", "interpolate").pack(fill="x", pady=3)

        tk.Label(frame, text="噪声过滤强度 (0=关闭, 1-5 推荐):").pack(anchor="w")
        self.noise_slider = tk.Scale(frame, from_=0, to=5, orient=tk.HORIZONTAL,
                                     variable=self.noise_filter_strength,
                                     tickinterval=1, length=150)
        self.noise_slider.pack(fill="x", pady=3)

        tk.Label(frame, text="颜色量化颜色数 (0=关闭):").pack(anchor="w")
        self.slider = tk.Scale(frame, from_=0, to=32, orient=tk.HORIZONTAL,
                               variable=self.max_color_number, tickinterval=4, length=150,
                               command=lambda _: self._align_slider_value())
        self.slider.pack(fill="x", pady=3)

        tk.Checkbutton(frame, text="启用SLIC超像素", variable=self.enable_slic).pack(anchor="w", pady=5)

        tk.Button(frame, text="选择单张图片", command=self.select_image).pack(fill="x", pady=5)
        tk.Button(frame, text="选择图片文件夹", command=self.select_folder).pack(fill="x", pady=5)
        tk.Button(frame, text="开始生成", command=self.start_processing).pack(fill="x", pady=15)

        self.status_label = tk.Label(frame, text="状态: 等待")
        self.status_label.pack()

        self.canvas_size = 600
        self.canvas = tk.Canvas(self.root, width=self.canvas_size, height=self.canvas_size, bg="#808080")
        self.canvas.pack(side="right", padx=10, pady=10)

        if not algo_loaded:
            self.status_label.config(text=f"算法模块加载失败: {algo_load_error}")

    def _align_slider_value(self):
        """保证颜色量化滑块值为2~32的偶数"""
        v = self.max_color_number.get()
        if v == 0:
            return
        v = max(2, min(32, v))
        aligned = (v // 2) * 2
        if aligned != self.max_color_number.get():
            self.max_color_number.set(aligned)

    def schedule_auto_update(self, *_):
        if self.update_job:
            self.root.after_cancel(self.update_job)
        self.update_job = self.root.after(300, self.auto_update_fields)

    def select_image(self):
        path = filedialog.askopenfilename(filetypes=[("图片文件", "*.png;*.jpg;*.jpeg;*.bmp")])
        if path:
            self.image_path = path
            self.folder_path = None
            self.status_label.config(text=f"已选择图片: {os.path.basename(path)}")
            img = Image.open(path)
            self.orig_w, self.orig_h = img.size
            self.auto_update_fields()
            self.show_preview(path)

    def select_folder(self):
        path = filedialog.askdirectory()
        if path:
            self.folder_path = path
            self.image_path = None
            self.status_label.config(text=f"已选择文件夹: {path}")
            for filename in os.listdir(path):
                if filename.lower().endswith((".png", ".jpg", ".jpeg", ".bmp")):
                    img_sample = Image.open(os.path.join(path, filename))
                    self.orig_w, self.orig_h = img_sample.size
                    break
            self.auto_update_fields()
            self.show_preview(None)

    def show_preview(self, img_path):
        """显示预览（支持文件夹模式下已处理的图片）"""
        self.canvas.delete("all")
        try:
            if img_path:
                img = Image.open(img_path)
            elif self.processed_images:
                img = self.processed_images[self.current_index]
            else:
                self.canvas.create_text(self.canvas_size // 2, self.canvas_size // 2,
                                        text="无预览", fill="white", font=("Arial", 28))
                return

            w, h = img.size
            scale = min(self.canvas_size / max(w, 1), self.canvas_size / max(h, 1))
            img = img.resize((int(w * scale), int(h * scale)), Image.NEAREST)
            self.tk_img = ImageTk.PhotoImage(img)
            self.canvas.create_image(self.canvas_size // 2, self.canvas_size // 2, image=self.tk_img)
        except Exception as e:
            self.canvas.create_text(self.canvas_size // 2, self.canvas_size // 2,
                                    text=f"预览错误\n{e}", fill="white", font=("Arial", 18))

    def show_prev_image(self, event=None):
        if self.processed_images:
            self.current_index = (self.current_index - 1) % len(self.processed_images)
            self.show_preview(None)

    def show_next_image(self, event=None):
        if self.processed_images:
            self.current_index = (self.current_index + 1) % len(self.processed_images)
            self.show_preview(None)

    def auto_update_fields(self):
        if not self.orig_w or not self.orig_h or self.updating_fields:
            return
        self.updating_fields = True
        try:
            tw = int(self.target_width_str.get()) if self.target_width_str.get().isdigit() else None
            th = int(self.target_height_str.get()) if self.target_height_str.get().isdigit() else None
            rs = int(self.region_size_str.get()) if self.region_size_str.get().isdigit() else None
            if tw:
                self.target_height_str.set(str(int(self.orig_h * (tw / self.orig_w))))
                self.region_size_str.set(str(max(1, self.orig_w // tw)))
            elif th:
                tw_calc = int(self.orig_w * (th / self.orig_h))
                self.target_width_str.set(str(tw_calc))
                self.region_size_str.set(str(max(1, self.orig_w // tw_calc)))
            elif rs:
                tw_calc = max(1, self.orig_w // rs)
                self.target_width_str.set(str(tw_calc))
                self.target_height_str.set(str(int(self.orig_h * (tw_calc / self.orig_w))))
        finally:
            self.updating_fields = False

    def start_processing(self):
        if not algo_loaded:
            messagebox.showerror("错误", "算法模块加载失败")
            return
        if not (self.image_path or self.folder_path):
            messagebox.showerror("错误", "请先选择图片或文件夹")
            return

        try:
            region_size = int(self.region_size_str.get())
            max_colors = int(self.max_color_number.get())
            max_colors_param = None if max_colors == 0 else max_colors
        except Exception as e:
            messagebox.showerror("错误", f"参数格式错误: {e}")
            return

        mode = self.compress_mode.get()
        enable_pixelization = (mode == "pixel")
        enable_compression = (mode == "interpolate")

        params = {
            "region_size": region_size,
            "color_weight": self.color_weight.get(),
            "iter_num": self.iter_num.get(),
            "enable_pixelization": enable_pixelization,
            "enable_color_quant": (max_colors_param is not None),
            "quant_colors": max_colors_param,
            "enable_slic": bool(self.enable_slic.get()),
            "enable_domain_interp": False,
            "enable_compression": enable_compression,
            "target_width": int(self.target_width_str.get()) if self.target_width_str.get().isdigit() else None,
            "target_height": int(self.target_height_str.get()) if self.target_height_str.get().isdigit() else None,
            "enable_noise_filter": self.noise_filter_strength.get()
        }

        self.status_label.config(text="处理中...")

        def worker():
            try:
                base_out_dir = os.path.join(os.getcwd(), "out")
                os.makedirs(base_out_dir, exist_ok=True)

                self.processed_images = []
                self.current_index = 0

                if self.image_path:
                    img = Image.open(self.image_path).convert("RGB")
                    results = slic_torch.run(img, **params)
                    save_path = os.path.join(base_out_dir, "output.png")
                    results[0].save(save_path)
                    self.processed_images.append(results[0])
                    self.root.after(0, self.show_preview, save_path)
                    self.root.after(0, lambda: self.status_label.config(text=f"完成：{save_path}"))
                else:
                    timestamp = time.strftime("%Y%m%d_%H%M%S")
                    folder_out_dir = os.path.join(base_out_dir, timestamp)
                    os.makedirs(folder_out_dir, exist_ok=True)

                    count = 0
                    for fn in os.listdir(self.folder_path):
                        if fn.lower().endswith((".png", ".jpg", ".jpeg", ".bmp")):
                            img = Image.open(os.path.join(self.folder_path, fn)).convert("RGB")
                            results = slic_torch.run(img, **params)
                            save_path = os.path.join(folder_out_dir, f"output_{fn}")
                            results[0].save(save_path)
                            self.processed_images.append(results[0])
                            count += 1
                            # ✅ 实时更新预览
                            self.root.after(0, self.show_preview, None)

                    self.root.after(0, lambda: self.status_label.config(
                        text=f"文件夹完成：{count} 张，方向键←/→切换预览，保存于 {folder_out_dir}"))
            except Exception as e:
                traceback.print_exc()
                self.root.after(0, lambda: self.status_label.config(text=f"处理失败: {e}"))
                messagebox.showerror("错误", f"处理失败: {e}")

        threading.Thread(target=worker, daemon=True).start()

if __name__ == "__main__":
    root = tk.Tk()
    app = PixelArtApp(root)
    root.mainloop()
